package org.hiucung.hlsdownloader.utils

import okio.buffer
import okio.sink
import okio.source
import java.io.File
import javax.crypto.Cipher
import javax.crypto.spec.IvParameterSpec
import javax.crypto.spec.SecretKeySpec

/**
 * @title: OkioAesDecrypt
 * @projectName HLS Downloader
 * @description: TS文件解密
 * @author Kicc
 * @date 2025/8/23 14:16
 */
class OkioAesDecrypt {

    /**
     * 解密文件 - 自动选择适合的方式（小文件用内存，大文件用流）
     */
    fun decryptFile(
        inputFile: File,
        outputFile: File,
        key: ByteArray,
        iv: ByteArray,
        useStreaming: Boolean? = null
    ): Boolean {
        return try {
            // 自动判断：大于5MB使用流式处理
            val shouldUseStreaming = useStreaming ?: (inputFile.length() > 5 * 1024 * 1024)

            if (shouldUseStreaming) {
                decryptFileStreaming(inputFile, outputFile, key, iv)
            } else {
                decryptFileInMemory(inputFile, outputFile, key, iv)
            }
        } catch (e: Exception) {
            e.printStackTrace()
            false
        }
    }

    /**
     * 内存方式解密（适合小文件）
     */
    private fun decryptFileInMemory(
        inputFile: File,
        outputFile: File,
        key: ByteArray,
        iv: ByteArray
    ): Boolean {
        return try {
            val cipher = createCipher(key, iv)

            inputFile.source().buffer().use { source ->
                outputFile.sink().buffer().use { sink ->
                    val encryptedData = source.readByteArray()
                    val decryptedData = cipher.doFinal(encryptedData)
                    sink.write(decryptedData)
                }
            }
            true
        } catch (e: Exception) {
            e.printStackTrace()
            false
        }
    }

    /**
     * 流式解密（适合大文件）
     */
    private fun decryptFileStreaming(
        inputFile: File,
        outputFile: File,
        key: ByteArray,
        iv: ByteArray
    ): Boolean {
        return try {
            val cipher = createCipher(key, iv)
            val buffer = ByteArray(64 * 1024) // 64KB缓冲区

            inputFile.source().buffer().use { source ->
                outputFile.sink().buffer().use { sink ->
                    var bytesRead: Int
                    while (source.read(buffer).also { bytesRead = it } != -1) {
                        val decryptedData = if (bytesRead == buffer.size) {
                            cipher.update(buffer)
                        } else {
                            cipher.doFinal(buffer, 0, bytesRead)
                        }
                        decryptedData?.let { sink.write(it) }
                    }
                }
            }
            true
        } catch (e: Exception) {
            e.printStackTrace()
            false
        }
    }

    /**
     * 创建AES解密器
     */
    private fun createCipher(key: ByteArray, iv: ByteArray): Cipher {
        val secretKey = SecretKeySpec(key, "AES")
        val ivSpec = IvParameterSpec(iv)
        val cipher = Cipher.getInstance("AES/CBC/PKCS7Padding")
        cipher.init(Cipher.DECRYPT_MODE, secretKey, ivSpec)
        return cipher
    }

    /**
     * 十六进制字符串转字节数组
     */
    fun hexStringToByteArray(hex: String): ByteArray {
        require(hex.length % 2 == 0) { "十六进制字符串长度必须为偶数" }
        return hex.chunked(2).map { it.toInt(16).toByte() }.toByteArray()
    }
}